#include "log.h"
#include "Share.h"
#include "ShareTransport.h"

#define MAX_RETRIES 5
#define RETRY_DELAY_MS 1000

ShareTransport::ShareTransport(std::function<HDEVNOTIFY(HANDLE)> registerHandleNotify)
    : m_bInitialised(FALSE),
      m_registerHandleNotify(registerHandleNotify),
      m_clipboardTransNotify(nullptr),
      m_fuseServerTransNotify(nullptr),
      m_fuseServer(GetFSL(), [](IFSL* Srv) { Srv->Release(); })
{ }

ShareTransport::~ShareTransport()
{
    if (m_clipboardTransNotify != nullptr) {
        UnregisterDeviceNotification(m_clipboardTransNotify);
        m_clipboardTransNotify = nullptr;
    }

    if (m_fuseServerTransNotify != nullptr) {
        UnregisterDeviceNotification(m_fuseServerTransNotify);
        m_fuseServerTransNotify = nullptr;
    }
}

DWORD ShareTransport::Init()
{
    if (m_bInitialised) {
        LOGE("ShareTransport already initialised!");
        return ERROR_ALREADY_INITIALIZED;
    }

    BOOL bSuccess = FALSE;
    for (INT id = 0; id < MAX_RETRIES; ++id) {
        if (!m_vios.InitDev(VIRTIO_SERIAL_PORT)) {
            Sleep(RETRY_DELAY_MS);
            continue;
        }

        if (!m_fuseServer->StartServer(FUSE_SERIAL_PORT)) {
            Sleep(RETRY_DELAY_MS);
            continue;
        }

        bSuccess = TRUE;
        break;
    }

    if (!bSuccess) {
        LOGE("Failed to init ShareTransport");
        return ERROR_CONNECTION_INVALID;
    }

    m_clipboardTransNotify = m_registerHandleNotify(GetClipboardTransportHandle());
    m_fuseServerTransNotify = m_registerHandleNotify(GetFuseServerTransportHandle());
    m_bInitialised = TRUE;
    return ERROR_SUCCESS;
}

HANDLE ShareTransport::GetClipboardTransportHandle()
{
    return m_vios.GetHandle();
}

HANDLE ShareTransport::GetFuseServerTransportHandle()
{
    if (m_fuseServer != nullptr) {
        return m_fuseServer->GetTransportHandle();
    }

    return INVALID_HANDLE_VALUE;
}

void ShareTransport::SetFuseServerReconnect(BOOL isReconnect)
{
    m_fuseServer->SetPeerReconnect(isReconnect);
}

BOOL ShareTransport::CancelIO()
{
    if (!m_bInitialised) {
        return FALSE;
    }

    return m_vios.PurgeIO();
}

BOOL ShareTransport::RecvHeader(ShareRecvProcess& recvProcess, DWORD &headerLen, SIZE_T &dataLen)
{
    if (!m_bInitialised) {
        return FALSE;
    }

    ShareHeader *header = nullptr;
    VioSOpStatus res = VioSOpStatus::SUCCESS;
    recvProcess.Clear();
    PBYTE headerBuf = recvProcess.Resize(sizeof(ShareHeader));
    if (headerBuf == nullptr) {
        LOGE("Failed to resize buffer for ShareHeader, requested size: {} bytes.", sizeof(ShareHeader));
        return FALSE;
    }

    do {
        res = m_vios.ReadAllTimed(headerBuf, sizeof(ShareHeader), RECV_DATA_TOTAL_DELAY);
        if (res == VioSOpStatus::SUCCESS) {
        } else if (res == VioSOpStatus::TIMEOUT) {
            continue;
        } else {
            LOGE("Failed to receive header data.");
            return FALSE;
        }

        header = reinterpret_cast<ShareHeader*>(headerBuf);
        if (!recvProcess.CheckHeader(header)) {
            LOGE("Clipboard header data checking error, discarding data.");
            DropLeftData();
            continue;
        }

        headerLen = header->headerLen;
        dataLen = header->dataLen;
        if (recvProcess.Resize(headerLen + dataLen) == nullptr) {
            LOGE("Failed to resize buffer, requested size: {} bytes.", headerLen + dataLen);
            return FALSE;
        }
        return TRUE;
    } while (1);
}

BOOL ShareTransport::RecvAllData(ShareRecvProcess& recvProcess)
{
    if (!m_bInitialised) {
        return FALSE;
    }

    VioSOpStatus res = VioSOpStatus::SUCCESS;

    do {
        DWORD headerLen = 0;
        SIZE_T dataLen = 0;
        if (!RecvHeader(recvProcess, headerLen, dataLen)) {
            LOGE("Failed to receive header data.");
            return FALSE;
        }

        PBYTE headerBuf = recvProcess.GetData();
        DWORD buffSize = static_cast<DWORD>(headerLen - sizeof(ShareHeader));
        if (buffSize) {
            res = m_vios.ReadAllTimed(reinterpret_cast<PBYTE>(headerBuf + sizeof(ShareHeader)),
                                      buffSize, RECV_DATA_TOTAL_DELAY);
            if (res == VioSOpStatus::SUCCESS) {
            } else if (res == VioSOpStatus::TIMEOUT) {
                LOGE("Failed to receive header within {} ms", RECV_DATA_TOTAL_DELAY);
                continue;
            } else {
                LOGE("Failed to receive extra header data.");
                return FALSE;
            }
        }

        if (!recvProcess.CheckExtraHeader()) {
            LOGE("Received invalid data.");
            DropLeftData();
            continue;
        }

        if (dataLen == 0) {
            if (recvProcess.GetFormat() == CBSH_SYNC_DATA_TO_HOST || recvProcess.GetFormat() == CBSH_FMT_URI_LIST) {
                return TRUE;
            }
            continue;
        }

        PBYTE dataBuf = reinterpret_cast<PBYTE>(headerBuf);
        res = m_vios.ReadAllTimed(dataBuf, static_cast<DWORD>(dataLen), RECV_DATA_TOTAL_DELAY);
        if (res == VioSOpStatus::SUCCESS) {
        } else if (res == VioSOpStatus::TIMEOUT) {
            LOGE("Failed to receive all data within {} ms", RECV_DATA_TOTAL_DELAY);
            continue;
        } else {
            LOGE("Failed to receive data with error: {}", (int)res);
            return FALSE;
        }

        return TRUE;
    } while (1);
}

BOOL ShareTransport::Send(PBYTE buf, DWORD bufLen)
{
    while (bufLen) {
        DWORD len = bufLen;
        if (!m_vios.Write(buf, len)) {
            LOGE("Send failed with error: {}", GetLastError());
            return FALSE;
        }
        bufLen -= len;
        buf += len;
    }
    return TRUE;
}

void ShareTransport::SendShareMsg(PBYTE hdr, SIZE_T hdrSize, ShareSendProcess* proc)
{
    if (!m_bInitialised) {
        return;
    }
    std::lock_guard<std::mutex> guard(m_mutex);

    if (!Send(hdr, static_cast<DWORD>(hdrSize))) {
        LOGE("Send header data error.");
        return;
    }

    if (!proc) {
        return;
    }

    PBYTE dataChunk = nullptr;
    DWORD chunkSize = 0;
    while (proc->NextData(&dataChunk, &chunkSize)) {
        if (!Send(dataChunk, chunkSize)) {
            LOGE("Send data error.");
            return;
        }
    }
}

void ShareTransport::DropLeftData()
{
    if (!m_bInitialised) {
        return;
    }
    m_vios.DropLeftData();
}

IFSL* ShareTransport::GetFuseServer()
{
    if (!m_bInitialised) {
        return nullptr;
    }
    return m_fuseServer.get();
}
